from tensorflow.keras import Model,layers
from tensorflow.keras.layers import Conv2D, Activation, BatchNormalization, Dense, Add, Concatenate
import tensorflow as tf
import numpy as np
from data.config import cfg

def mish(x):
    return x * tf.math.tanh(tf.math.softplus(x))
    # return tf.keras.layers.Lambda(lambda x: x*tf.tanh(tf.math.log(1+tf.exp(x))))(x)

# CBM模块
class ConvBNMish(layers.Layer):
    def __init__(self, filters, kernel_size, strides=1, padding='same' ):
        super(ConvBNMish, self).__init__()
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding

    def build(self, input_shape):
        self.conv = Conv2D(self.filters, self.kernel_size, strides=self.strides, padding=self.padding)
        self.bn = BatchNormalization()
        self.mish = mish

    def call(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.mish(x)
        return x

# CBM模块
class ConvBNLeak(layers.Layer):
    def __init__(self, filters, kernel_size, strides=1, padding='same'):
        super(ConvBNLeak).__init__()
        self.filters = filters
        self.kernel_size = kernel_size
        self.strides = strides
        self.padding = padding

    def build(self, input_shape):
        self.conv = Conv2D(self.filters, self.kernel_size, strides=self.strides, padding=self.padding)
        self.bn = BatchNormalization()
        self.relu = Activation()

    def call(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x


# CSPX模块
class CSPBlock(layers.Layer):
    def __init__(self, filters, block_num):
        super(CSPBlock).__init__()
        self.filters = filters
        self.block_num = block_num

    def build(self, input_shape):
        # 先设置步长为2，下采样 图片缩小一半
        self.cbm = ConvBNMish(self.filters, 3, strides=2)
        self.cbm_1 = ConvBNMish(self.filters, 1)
        self.cbm_kernel_1 = ConvBNMish(self.filters//2, 1,)
        self.cbm_kernel_3 = ConvBNMish(self.filters//2, 3,)

    def call(self, x):
        x = self.cbm(x)
        # 残差边
        res_block = self.cbm_kernel_1(x)

        # 主模块
        x = self.cbm_kernel_1(x)
        for i in range(self.block_num):
            y = self.cbm_kernel_1(x)
            y = self.cbm_kernel_3(y)
            x = Add()([x, y])
        x = self.cbm_kernel_1(x)

        x = Concatenate()([res_block, x])

        x = self.cbm_1(x)
        return x


class CSP_darknet(layers.Layer):
    def __init__(self):
        super(CSP_darknet, self).__init__()

    def build(self, input_shape):
        self.CBM = ConvBNMish(32, 3, strides=1, padding='same')
        self.csp1 = CSPBlock(64, 1)
        self.csp2 = CSPBlock(128, 2)
        self.csp8_1 = CSPBlock(256, 8)
        self.csp8_2 = CSPBlock(512, 8)
        self.csp4 = CSPBlock(1024, 4)


    def call(self, x):
       x = self.CBM(x)
       x = self.csp1(x)
       x = self.csp2(x)
       x = self.csp8_1(x)
       feature1 = x
       x = self.csp8_2(x)
       feature2 = x
       x = self.csp4(x)
       feature3 = x

       return feature1, feature2, feature3




class YOLOV4(Model):
    def __init__(self):
        super(YOLOV4, self).__init__()
        self.cspdarknet = CSP_darknet()

    def call(self, x):
        x = self.cspdarknet(x)
        return x

if __name__ == '__main__':
    model = YOLOV4()
    image = np.full((1, 416, 416, 3), fill_value=128.0)
    x = model(image)
    model.summary()
    print(x.shape)
    pass